package pku;

import java.io.BufferedOutputStream;
import java.io.FileOutputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

public class Producer
{
    private static final int CACHE_SIZE = 20000;
    private static final Map<String, BufferedOutputStream> topicStreams = new HashMap<>();

    private final Map<String, List<ByteMessage>> msgs = new HashMap<>();
    private int counter = 0;

    private void store(BufferedOutputStream bw, List<ByteMessage> messages) throws Exception
    {
        for (ByteMessage b : messages)
        {
            byte[] header = Util.headerToString(b.headers(), ",").getBytes();
            byte[] body = b.getBody();

            bw.write(Util.toBytes(header.length));
            bw.write(header);
            bw.write(Util.toBytes(body.length));
            bw.write(body);
        }
        bw.flush();
    }

    private void deferFlush() throws Exception
    {
        if (counter > CACHE_SIZE)
        {
            counter = 0;
            flush();
        }
        counter++;
    }

    public ByteMessage createBytesMessageToTopic(String topic, byte[] body)
    {
        ByteMessage msg = new DefaultMessage(body);
        msg.putHeaders(MessageHeader.TOPIC, topic);
        return msg;
    }

    public void send(ByteMessage msg) throws Exception
    {
        String topic = msg.headers().getString(MessageHeader.TOPIC);
        if (!msgs.containsKey(topic))
        {
            msgs.put(topic, new LinkedList<>());
        }
        msgs.get(topic).add(msg);
        deferFlush();
    }

    public void flush() throws Exception
    {
        synchronized (topicStreams)
        {
            for (String topic : msgs.keySet())
            {
                if (!topicStreams.containsKey(topic))
                {
                    OutputStream out = new FileOutputStream(Util.DATA_PATH + topic, true);
                    BufferedOutputStream buffer = new BufferedOutputStream(out);
                    topicStreams.put(topic, buffer);
                }
                BufferedOutputStream bw = topicStreams.get(topic);

                List<ByteMessage> messages = msgs.get(topic);
                store(bw, messages);
                messages.clear();
            }
        }
    }
}
